import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

def plot_fig5():
    csvs = [
    ['dataset/third_time_images/easy.csv', 'easy'],
    ['dataset/third_time_images/medium.csv', 'medium'],
    ['dataset/third_time_images/hard.csv', 'hard'],
    ]

    rename_dict = {
        'DGN_R2': 'DGN',
        'TranDGN': 'TMAC(ours)',
        'GACPN_R3': 'MAGIC',
    }

    sns.set_style('whitegrid')
    fig, axs = plt.subplots(1, 3, figsize=(15, 5))

    for i ,csv in enumerate(csvs):
        path, diff = csv
        data = pd.read_csv(path)
        data.insert(data.shape[1], 'diff', diff)
        data = data.replace(rename_dict.keys(), rename_dict.values())

        axes = sns.lineplot(ax=axs[i], data=data, x='time_stamp', y='score', hue='model_type')
        axes.set_xlabel("Episode")
        axes.set_ylabel("Reward")
        legend = axes.get_legend()
        legend.set_loc('lower right')
        legend.set_title('')
        legend.set_bbox_to_anchor((1.0, 0.0))

    plt.tight_layout(rect=[0, 0.08, 1, 1])
    label_font = {'family': 'Times New Roman', 'size': 18, 'weight': 'bold'}
    fig.text(0.18, 0.02, f"(a) Easy environment", ha='center', fontdict=label_font)
    fig.text(0.52, 0.02, f"(b) Medium environment", ha='center', fontdict=label_font)
    fig.text(0.85, 0.02, f"(c) Hard environment", ha='center', fontdict=label_font)
    fig.savefig('res.pdf')

def plot_fig6():
    csvs = [
        ['dataset/smac/25m.csv', '25m'],
        ['dataset/smac/mmm2.csv', 'mmm2']
    ]

    rename_dict = {}
    hue_order = ['DGN', 'TMAC', 'MAGIC']

    sns.set_style('whitegrid')
    fig, axs = plt.subplots(1, 2, figsize=(10, 5))

    for i ,csv in enumerate(csvs):
        path, diff = csv
        data = pd.read_csv(path)
        data.insert(data.shape[1], 'diff', diff)
        data = data.replace(rename_dict.keys(), rename_dict.values())

        axes = sns.lineplot(ax=axs[i], data=data, x='timesteps', y='scores', hue='types', hue_order=hue_order)
        axes.set_xlabel("Timesteps")
        axes.set_ylabel("Reward")
        legend = axes.get_legend()
        legend.set_loc('lower right')
        legend.set_title('')
        legend.set_bbox_to_anchor((1.0, 0.0))

    plt.tight_layout(rect=[0, 0.08, 1, 1])
    label_font = {'family': 'Times New Roman', 'size': 18, 'weight': 'bold'}
    fig.text(0.28, 0.02, f"(a) 25m", ha='center', fontdict=label_font)
    fig.text(0.76, 0.02, f"(b) MMM2", ha='center', fontdict=label_font)
    fig.savefig('res.pdf')

def plot_fig7():
    csvs = [
        ['dataset/multi_obs-multi/easy-multi.csv', 'easy'],
        ['dataset/multi_obs-multi/medium-multi.csv', 'medium'],
        ['dataset/multi_obs-multi/hard-multi.csv', 'hard'],
    ]

    rename_dict = {
        # 'DGN_R2': 'DGN',
        'TranDGN': 'TMAC',
        'TranDGN-multi': 'TMAC w/o multi',
    }

    hue_order = ['TMAC w/o multi', 'TMAC']

    sns.set_style('whitegrid')
    fig, axs = plt.subplots(1, 3, figsize=(15, 5))

    for i ,csv in enumerate(csvs):
        path, diff = csv
        data = pd.read_csv(path)
        data.insert(data.shape[1], 'diff', diff)
        data = data.replace(rename_dict.keys(), rename_dict.values())

        axes = sns.lineplot(ax=axs[i], data=data, x='time_stamp', y='score', hue='model_type', hue_order=hue_order)
        axes.set_xlabel("Episode")
        axes.set_ylabel("Reward")
        legend = axes.get_legend()
        legend.set_loc('lower right')
        legend.set_title('')
        legend.set_bbox_to_anchor((1.0, 0.0))

    plt.tight_layout(rect=[0, 0.08, 1, 1])
    label_font = {'family': 'Times New Roman', 'size': 18, 'weight': 'bold'}
    fig.text(0.18, 0.02, f"(a) Easy environment", ha='center', fontdict=label_font)
    fig.text(0.52, 0.02, f"(b) Medium environment", ha='center', fontdict=label_font)
    fig.text(0.85, 0.02, f"(c) Hard environment", ha='center', fontdict=label_font)
    fig.savefig('res.pdf')

def plot_fig8():
    csvs = [
        ['dataset/multi_obs-processor/easy_processor.csv', 'easy'],
        ['dataset/multi_obs-processor/medium-processor.csv', 'medium'],
        ['dataset/multi_obs-processor/hard-processor.csv', 'hard'],
    ]

    rename_dict = {
        # 'DGN_R2': 'DGN',
        'TranDGN': 'TMAC',
        'TranDGN-processor': 'TMAC w/o processor',
    }

    hue_order = ['TMAC w/o processor', 'TMAC']

    sns.set_style('whitegrid')
    fig, axs = plt.subplots(1, 3, figsize=(15, 5))

    for i ,csv in enumerate(csvs):
        path, diff = csv
        data = pd.read_csv(path)
        data.insert(data.shape[1], 'diff', diff)
        data = data.replace(rename_dict.keys(), rename_dict.values())

        axes = sns.lineplot(ax=axs[i], data=data, x='time_stamp', y='score', hue='model_type', hue_order=hue_order)
        axes.set_xlabel("Episode")
        axes.set_ylabel("Reward")
        legend = axes.get_legend()
        legend.set_loc('lower right')
        legend.set_title('')
        legend.set_bbox_to_anchor((1.0, 0.0))

    plt.tight_layout(rect=[0, 0.08, 1, 1])
    label_font = {'family': 'Times New Roman', 'size': 18, 'weight': 'bold'}
    fig.text(0.18, 0.02, f"(a) Easy environment", ha='center', fontdict=label_font)
    fig.text(0.52, 0.02, f"(b) Medium environment", ha='center', fontdict=label_font)
    fig.text(0.85, 0.02, f"(c) Hard environment", ha='center', fontdict=label_font)
    fig.savefig('res.pdf')

def plot_fig9():
    csvs = [
        ['dataset/with MLP/easy_with_MLP.csv', 'easy'],
        ['dataset/with MLP/medium_with_MLP.csv', 'medium'],
        ['dataset/with MLP/hard_with_MLP.csv', 'hard'],
    ]

    rename_dict = {}

    hue_order = ['TMAC w/ MLP', 'TMAC']

    sns.set_style('whitegrid')
    fig, axs = plt.subplots(1, 3, figsize=(15, 5))

    for i ,csv in enumerate(csvs):
        path, diff = csv
        data = pd.read_csv(path)
        data.insert(data.shape[1], 'diff', diff)
        data = data.replace(rename_dict.keys(), rename_dict.values())

        axes = sns.lineplot(ax=axs[i], data=data, x='episodes', y='scores', hue='types', hue_order=hue_order)
        axes.set_xlabel("Episode")
        axes.set_ylabel("Reward")
        legend = axes.get_legend()
        legend.set_loc('lower right')
        legend.set_title('')
        legend.set_bbox_to_anchor((1.0, 0.0))

    plt.tight_layout(rect=[0, 0.08, 1, 1])
    label_font = {'family': 'Times New Roman', 'size': 18, 'weight': 'bold'}
    fig.text(0.18, 0.02, f"(a) Easy environment", ha='center', fontdict=label_font)
    fig.text(0.52, 0.02, f"(b) Medium environment", ha='center', fontdict=label_font)
    fig.text(0.85, 0.02, f"(c) Hard environment", ha='center', fontdict=label_font)
    fig.savefig('res.pdf')

if __name__ == '__main__':
    plot_fig6()